import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
from PIL import ImageGrab
import keyboard
import os
import time

# 英雄头像（同时也可用来存储星级图案）
class Hero:
    # 从底部去掉的范围
    outRange = 3
    # 英雄数量
    __number = 0
    # 识别阈值
    __threshold = 0.65

    def __init__(self, image, name):
        self.__name = name
        self.__image = self.__preConvertImage(image)
    # 预处理图像，将透明部分替换为白色

    def __preConvertImage(self, image):
        # 处理图像Alpha通道，将透明部分替换为不透明黑色背景
        for i in image:
            for j in i:
                if j[3] == 0:
                    j[0:3] = [0,0,0]
        return cv.cvtColor(image, cv.COLOR_BGRA2BGR)

    def returnImage(self):
        return self.__image

    def show(self):
        cv.imshow(self.__class__.__name__, self.__image)
        cv.waitKey(0)
        cv.destroyAllWindows()

    def changeNumber(self, number):
        self.__number = number

    def checkNumber(self):
        return self.__number

    def checkName(self):
        return self.__name

    def changeThreshold(self, number):
        self.__threshold = number

    def checkThreshold(self):
        return self.__threshold

class BackGround:
    w, h = [32, 32]
    xstart, xend = [430, 1650]
    ystart, yend = [135, 950]
    star1 = Hero(cv.imread('images/star1.png', cv.IMREAD_UNCHANGED ), 'star1')
    star2 = Hero(cv.imread('images/star2.png', cv.IMREAD_UNCHANGED ), 'star2')
    star3 = Hero(cv.imread('images/star3.png', cv.IMREAD_UNCHANGED ), 'star3')
    def __init__(self, image):
        self.__image = image

    def match(self, hero):
        temp = hero.returnImage()
        matchImage = np.zeros([12,12,3], dtype = np.uint8)
        # 提高准确率，只匹配小图标中间的一部分
        for i in range(12):
            for j in range(12):
                matchImage[i][j] = temp[i + 10][j + 10]

        # 匹配英雄头像
        res = cv.matchTemplate(self.__image[self.ystart:self.yend, self.xstart:self.xend], matchImage, cv.TM_CCOEFF_NORMED)
        loc = np.where(res >= hero.checkThreshold())
        num = 0
        prePt = None
        for pt in zip(*loc[::-1]):
            # 去除模板匹配因误差而造成重复的对象
            if prePt != None and abs(prePt[0]-pt[0]) < self.w and abs(prePt[1]-pt[1]) < self.h:
                prePt = pt
                continue

            color = ()
            # 按星级匹配对象
            if (cv.matchTemplate(self.__image[self.ystart+pt[1]+15:self.ystart+pt[1]+35, self.xstart+pt[0]-15:self.xstart+pt[0]+25], self.star3.returnImage(), cv.TM_CCOEFF_NORMED) > 0.8).any():
                num += 9
                color = (0, 0, 255)
            elif (cv.matchTemplate(self.__image[self.ystart+pt[1]+15:self.ystart+pt[1]+35, self.xstart+pt[0]-5:self.xstart+pt[0]+20], self.star2.returnImage(), cv.TM_CCOEFF_NORMED) > 0.6).any():
                num += 3
                color = (0, 255, 0)
            elif (cv.matchTemplate(self.__image[self.ystart+pt[1]+15:self.ystart+pt[1]+35, self.xstart+pt[0]:self.xstart+pt[0]+15], self.star1.returnImage(), cv.TM_CCOEFF_NORMED) > 0.5).any():
                num += 1
                color = (255, 0, 0)
            prePt = pt
            # 将匹配结果框选
            if color != ():
                cv.rectangle(self.__image, (self.xstart+pt[0]-10, self.ystart+pt[1]-10), (self.xstart+pt[0] + self.w-10, self.ystart+pt[1] + self.h-10), color, 2)
        hero.changeNumber(num)

    def show(self):
        temp = cv.resize(self.__image,None,fx=0.75, fy=0.75, interpolation = cv.INTER_CUBIC)
        cv.imshow(self.__class__.__name__, temp)
        cv.waitKey(5000)
        cv.destroyAllWindows()

def changeThreshold(name, threshold):
    global heros, heroName
    heros[heroName.index(name)].changeThreshold(threshold)

# 此为辅助器相关设置

heroName = ["Abaddon","Alchemist","Anti-Mage","Arc_Warden","Axe","Batrider","Beastmaster","Bloodseeker",
            "Bounty_Hunter","Chaos_Knight","Clockwerk","Crystal_Maiden","Disruptor","Doom","Dragon_Knight",
            "Drow_Ranger","Enchantress","Enigma","Gyrocopter","Juggernaut","Keeper_of_the_Light","Kunkka",
            "Lich","Lina","Lone_Druid","Luna","Lycan","Medusa","Mirana","Morphling","Naga_Siren","Nature's_Prophet",
            "Necrophos","Nyx_Assassin","Ogre_Magi","Omniknight","Phantom_Assassin","Puck","Pudge","Queen_of_Pain",
            "Razor","Sand_King","Shadow_Fiend","Shadow_Shaman","Slardar","Slark","Sniper","Techies",
            "Templar_Assassin","Terrorblade","Tidehunter","Timbersaw","Tinker","Tiny","Treant_Protector",
            "Troll_Warlord","Tusk","Venomancer","Viper","Warlock","Windranger","Witch_Doctor"]
heroNumber = 62
heros = []

for name in heroName:
    heros.append(Hero(cv.imread('images/'+name+'_minimap_icon.png', cv.IMREAD_UNCHANGED ), name))

# 手动调整阈值
changeThreshold("Gyrocopter", 0.8)
changeThreshold("Terrorblade", 0.5)
changeThreshold("Windranger", 0.7)
changeThreshold("Necrophos", 0.5)
changeThreshold("Queen_of_Pain", 0.8)
changeThreshold("Treant_Protector", 0.5)
changeThreshold("Arc_Warden", 0.8)
changeThreshold("Dragon_Knight", 0.8)
changeThreshold("Shadow_Fiend", 0.75)
changeThreshold("Anti-Mage", 0.75)
changeThreshold("Witch_Doctor", 0.6)

if __name__ == '__main__':
    while True:
        keyboard.wait('shift+ctrl+alt+s')
        os.system('cls')
        print('..............start grab..................')
        # 由于显示比例不为100%时，会使ImageGrab截图不全，此处模拟按下screen print键截图
        keyboard.press('print screen')
        # 延迟一秒执行，防止截图到不了剪切板
        time.sleep(1)
        backGround = ImageGrab.grabclipboard()
        backGround = cv.cvtColor(np.asarray(backGround), cv.COLOR_RGB2BGR)
        backGround = BackGround(backGround)
        for hero in heros:
            hero.changeNumber(0)
            backGround.match(hero)
            print(hero.checkName() + ': ' + str(hero.checkNumber()))
        backGround.show()
